/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: Hccl
 * Author: baochangjin
 * Create: 2024-02-28
 */

#ifndef HCCL_RPC_NPU_H
#define HCCL_RPC_NPU_H

#include <map>
#include <mutex>

#include "hccl_comm_pub.h"
#include "hccl_comm_rpc.h"

#include "remote_access.h"
#include "../../hcom/src/hcom_private.h"

#ifdef __cplusplus
extern "C" {
#endif

HcclResult HcclFinalizeCsComm(HcclComm comm);
HcclResult HcclRpcRegisterGlobalMemory(void* addr, u64 size);
HcclResult HcclRpcUnregisterGlobalMemory(void* addr);
HcclResult HcomSetWorkspaceResource(const std::string &tag, const char *group, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize);

HcclResult HRpcCollRemotePairedParaCheck(const HcomRemoteOperationParams &params);

HcclResult HcomRemoteLookup(s32 tag, s32 insertFlag, void *stream, s32 &intZerocpyFlag, s32 &outZerocpyFlag,
    RemoteLookupPubParams &lookupPubInfo, u64 keyNum);

HcclResult HcomCollRemoteUpdate(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertOption,
    s32 valueItemSize, const char *group, void *stream, s32 maxEmbeddingDim, u64 keyNum);
HcclResult HcomCollRemoteLookup(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertFlag,
    s32 valueItemSize, const char *group, void *stream, s32 flags, s32 maxEmbeddingDim, u64 keyNum);

HcclResult HcomCollRemoteLookupPaired(const HcomRemoteOperationParams &params);
HcclResult HcomCollRemoteLookupUniquedAndPaired(const HcomRemoteOperationParams &params);
HcclResult HcomCollRemoteUpdatePaired(const HcomRemoteOperationParams &params);

HcclResult LoadHRpcCsInitTask(const char* rankTableM, s32 rankId);
HcclResult HcclRPCInitComm(const char* rankTableM, uint32_t rank, const CommAttr* attr,
    HcclComm* comm, HcomInfo &rankInfo);

HcclResult HcclWaitSomeWithLog(int32_t count, HcclRequest requestArray[], s32* compCount, s32 compIndices[],
    HcclStatus compStatus[], bool needRunLog);

int HcclRPCImprobe(int srcRank, int tag, HcclComm comm, int* flag, HcclMessage* msg, HcclStatus* status);
int HcclRPCGetCount(const HcclStatus *status, HcclDataType dataType, int *count);
int HcclRPCImrecv(void* buffer, int count, HcclDataType dataType, HcclMessage* msg, HcclRequest* request);
#ifdef __cplusplus
}
#endif
#endif  // HCCL_RPC_NPU_H

